
#include "sync_bn_training_update_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
    const uint32_t BLOCK_SIZE = 32;
    const uint32_t BUFFER_NUM  = 2;
    static ge::graphStatus TilingFunc(gert::TilingContext* context)
    {
        SyncBNTrainingUpdateCustomTilingData tiling;

        // 从算子属性中获取momentum和epsilon参数
        float momentum = 0.9f; // 默认值
        float epsilon = 1e-5f; // 默认值
        // 尝试从属性中获取实际值
        const gert::RuntimeAttrs *attrs = context->GetAttrs();
        if (attrs != nullptr) {
            const size_t momentum_index = 0;
            const size_t epsilon_index = 1;
        // 使用索引获取float属性值
            const float* momentum_ptr = attrs->GetFloat(momentum_index);
            if (momentum_ptr != nullptr) {
                momentum = *momentum_ptr;
            }
        
        // 使用索引获取float属性值
            const float* epsilon_ptr = attrs->GetFloat(epsilon_index);
            if (epsilon_ptr != nullptr) {
                epsilon = *epsilon_ptr;
            }
        }
        // 获取硬件信息
        uint64_t ubSize;// Unified Buffer(UB)大小
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);// 获取UB大小
        auto coreNum = ascendcPlatform.GetCoreNum();// 获取核心数量

        // 计算输入数据相关信息
        uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();// 输入元素数量
        uint32_t typeLength = 4;//float类型固定长度为4
        // ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);// 数据类型长度(字节)
        uint32_t inputLength = inputNum * typeLength;// 输入数据总字节数
        // uint32_t inputBytes = inputLength / inputNum;// 每个元素的字节数

        // 计算UB中可以存放的数据量
        uint32_t ubDataNumber = 6;// 根据数据类型确定UB中需要存放的数据种类数量
        // 计算每个数据可用的32B块数
        uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
        //计算每个tile可以处理的元素数量
        uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / typeLength;

         // 计算核心间数据分配
        uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);// 32B对齐后的输入长度
        coreNum = (coreNum <  inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;// 调整核心数
        coreNum = (coreNum >= 1) ? coreNum : 1;// 确保至少一个核心
        uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE / coreNum;// 每个核心处理的32B块数
        uint32_t tailBlockNum = (inputLengthAlgin32 / BLOCK_SIZE) % coreNum;// 剩余的不能均匀分配的块数
    
        // 计算小核心(处理较少数据的核心)的参数
        uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / typeLength;
        uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
        uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
        uint32_t smallTailDataNum = smallCoreDataNum - (tileDataNum * smallTileNum);
        smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;
    
        // 计算大核心(处理较多数据的核心)的参数
        everyCoreInputBlockNum += 1;// 大核心多处理一个32B块
        uint32_t bigCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / typeLength;
        uint32_t bigTileNum = everyCoreInputBlockNum / tileBlockNum;
        uint32_t finalBigTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? bigTileNum : bigTileNum + 1;
        uint32_t bigTailDataNum = bigCoreDataNum - tileDataNum * bigTileNum;
        bigTailDataNum = bigTailDataNum == 0 ? tileDataNum : bigTailDataNum;

        // 设置tiling参数
        tiling.set_smallCoreDataNum(smallCoreDataNum);
        tiling.set_bigCoreDataNum(bigCoreDataNum);
        tiling.set_tileDataNum(tileDataNum);
        tiling.set_smallTailDataNum(smallTailDataNum);
        tiling.set_bigTailDataNum(bigTailDataNum);
        tiling.set_finalSmallTileNum(finalSmallTileNum);
        tiling.set_finalBigTileNum(finalBigTileNum);
        tiling.set_tailBlockNum(tailBlockNum);
        tiling.set_momentum(momentum);
        tiling.set_epsilon(epsilon);
        
        // 设置执行参数
        context->SetBlockDim(coreNum);// 设置使用的核心数量
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); // 保存tiling数据
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());// 设置数据大小
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;// 不需要额外workspace空间
        return ge::GRAPH_SUCCESS;
    }
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
const auto inputDataType = context->GetInputDataType(0);
context->SetOutputDataType(0, inputDataType);
return ge::GRAPH_SUCCESS;
}
}


namespace ops {
class SyncBNTrainingUpdateCustom : public OpDef {
public:
    explicit SyncBNTrainingUpdateCustom(const char* name) : OpDef(name)
    {
        this->Input("local_mean")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("local_var")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("global_mean")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("global_var")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->Output("updated_mean")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("updated_var")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->Attr("momentum")
            .AttrType(OPTIONAL)
            .Float(0.9f); // 默认值0.9
        this->Attr("epsilon")
            .AttrType(OPTIONAL)
            .Float(1e-5f); // 默认值1e-5

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(SyncBNTrainingUpdateCustom);
}
